from torch import nn
import torch

__all__ = ['color_space_net']

from timm.models import vit_base_patch16_224
from timm.models.color_space import UnetColorSpaceGeneratorV2


class ColorSpaceNet(nn.Module):

    def __init__(self, in_chans, num_classes, pretrained=False):
        super(ColorSpaceNet, self).__init__()
        # self.generator = UnetColorSpaceGenerator(in_chans, 3)
        self.generator = UnetColorSpaceGeneratorV2(in_chans, 3)
        # self.generator = SiblingBranch(in_chans, 3)
        # self.generator = Self_Attn(3)
        # self.features = tf_efficientnet_b0_ns(pretrained=pretrained, num_classes=num_classes)
        self.features = vit_base_patch16_224(pretrained=pretrained, num_classes=num_classes)
        # self.features = resnet18(pretrained=pretrained, num_classes=1000)

    # def forward(self, x1, x2):
    def forward(self, x):
        x2 = self.generator(x)

        # output = []
        # output.append(x)
        # output.append(x2)
        # for i, feat in enumerate(output):
        #     print("第{}个:".format(i + 1), feat.shape)
        x3 = self.features(x2)
        return x3


def color_space_net(in_chans=3, num_classes=2, pretrained=True):
    model = ColorSpaceNet(in_chans=in_chans, num_classes=num_classes, pretrained=pretrained)
    model.default_cfg = model.features.default_cfg
    return model

if __name__ == "__main__":
    model = ColorSpaceNet(in_chans=3, num_classes=2, pretrained=True)
    x1 = torch.randn(4, 3, 224, 224)
    # x2 = torch.randn(4, 3, 224, 224)
    regression = model(x1)
